#ifndef TEST_CNN_CKKS_MNIST_HELPER
#define TEST_CNN_CKKS_MNIST_HELPER

#include "CNN/mnist_helper.hpp"
#include "linear/linear_hmm_engine.hpp"
#include "powerof2/linear_p2_crypto_agent.hpp"

static constexpr unsigned long MATRIX_DIM_M = 64l;
static constexpr unsigned long MATRIX_DIM_N = 64l;
static constexpr unsigned long MATRIX_DIM_K = 64l;
static constexpr unsigned long PLAINTEXT_NSLOT =
    MAX(MATRIX_DIM_M, MATRIX_DIM_K) * MAX(MATRIX_DIM_K, MATRIX_DIM_N);

class ckks_hmnist_client : public h_mnist_client<helib::CKKS> {
private:
  ckks_hmnist_client() = delete;

  helib::Ctxt init_hx(const torch::Tensor &a, int asz0) const;

  helib::Ctxt init_hconv(const torch::Tensor &a, int asz0, int asz1) const;

  helib::Ctxt init_hfc1(const torch::Tensor &a, int asz0) const;

  helib::Ctxt init_hfc2(const torch::Tensor &a) const;

public:
  ckks_hmnist_client(const params<helib::CKKS> &params);
  ~ckks_hmnist_client() { delete hmmcc_; }

  const std::shared_ptr<helib::Context> get_context() const {
    return hmmcc_->context();
  }

  // Generate encrypted model
  void init_hcnn(encrypted_model<helib::CKKS> &ctxt_model,
                 const torch::nn::Conv2d &conv1, const torch::nn::Linear &fc1,
                 const torch::nn::Linear &fc2) const override;

  // Generate encrypted image
  void init_input(encrypted_input<helib::CKKS> &ctxt_input,
                  const torch::Tensor &filter,
                  const torch::Tensor &x) const override;

  void recover_result(at::Tensor &out,
                      const encrypted_output<helib::CKKS> &in) const override;

private:
  ENC key = ENC::PUBKEY;

  linearp2_hmmcc<helib::CKKS> *hmmcc_;
};

class ckks_hmnist_server : public h_mnist_server<helib::CKKS> {
private:
  struct ckks_metric : public metric {
    double err = 0.0;
    double init_err = 0.0;
    double conv_hmul_err = 0.0;
    double conv_hadd_err = 0.0;
    double sq1_err = 0.0;
    double fc1_err = 0.0;
    double sq2_err = 0.0;
  };

  ckks_hmnist_server() = delete;

public:
  ckks_hmnist_server(const std::shared_ptr<helib::Context> context,
                     const hmm_status<shmm_engine> &status);
  ~ckks_hmnist_server() { delete hmme_; }

  void
  forward_test(encrypted_output<helib::CKKS> &out,
               const encrypted_input<helib::CKKS> &in,
               const encrypted_model<helib::CKKS> &ctxt_model) const override;

  void report_metrics(std::ostream &ss = std::cout) const override;

private:
  ckks_metric metric_;

  linear_hmme *hmme_;
};

#endif /* TEST_CNN_CKKS_MNIST_HELPER */
